/*
 * Copyright 2015 MingxingWang

 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at

 *     http://www.apache.org/licenses/LICENSE-2.0

 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.graphx.graphdb

import java.io._

import com.audaque.graphdb._
import com.audaque.graphdb.util.GraphDBUtil
import org.apache.spark.graphx.{EdgeDirection, VertexId}

import scala.collection.immutable.Set
import scala.collection.mutable

/**
 * Created by Ming on 2015/3/26.
 */

class LEdges(var partitions: List[LEdgePartition[Property, Property]] = List[LEdgePartition[Property, Property]](),
             val addBuffer: mutable.Map[VertexId, mutable.Map[VertexId, Property]] = new mutable.HashMap(),
             var addBuffSize: Int = 0) {

  // reversed data
  var reversEdges: LEdges = null

  def count(): Long = {
    val pSize = {
      val size = partitions.map(_.size)
      if (size.isEmpty) 0 else size.reduce(_ + _)
    }
    val total = pSize + addBuffSize
    println("Total:" + total + ", partition:" + pSize + ", add buffer:" + addBuffSize)
    total
  }

  def contain(srcId: VertexId, dstId: VertexId): Boolean = {
    find(srcId, dstId) match {
      case Some(x) => true
      case None => false
    }
  }

  def find(srcId: VertexId, dstId: VertexId): Option[Edge] = {
    addBuffer.get(srcId) match {
      case Some(dst) =>
        dst.get(dstId) match {
          case Some(x) =>
            return Some(Edge(srcId, dstId, x))
          case None =>
        }
      case None =>
    }
    for (p <- partitions) {
      p.find(srcId, dstId) match {
        case Some(ed) => return Some(Edge(srcId, dstId, ed))
        case None =>
      }
    }
    None
  }


  /**
   *
   * @param srcId
   * @param dstId
   * @param prop
   * @param replace : 冲突时是否替换原值
   * @return
   */
  def add(srcId: VertexId, dstId: VertexId, prop: Property, replace: Boolean = false): Boolean = {
    if (reversEdges != null) {
      reversEdges.add(dstId, srcId, prop)
    }
    addBuffer.get(srcId) match {
      case Some(dst) =>
        dst.get(dstId) match {
          case Some(x) =>
            if (replace) {
              dst += (dstId -> prop)
              return true
            } else {
              return false
            }
          case None =>
        }
      case None =>
    }
    if (replace) {
      for (p <- partitions) {
        if (p.update(srcId, dstId, prop)) return true
      }
    } else {
      for (p <- partitions) {
        if (p.isDefined(srcId, dstId)) return false
      }
    }
    addBuffer.get(srcId) match {
      case Some(x) =>
        x += (dstId -> prop)
      case None =>
        val x = mutable.Map[VertexId, Property](dstId -> prop)
        addBuffer += (srcId -> x)
    }
    addBuffSize += 1
    if (addBuffSize >= GraphDBConf.EdgeBufferMaxSize)
      update
    true
  }

  def modify(srcId: VertexId, dstId: VertexId, prop: Property): Boolean = {
    var res: Boolean = false
    if (reversEdges != null) {
      reversEdges.modify(dstId, srcId, prop)
    }
    addBuffer.get(srcId) match {
      case Some(dsts) =>
        if (dsts.contains(dstId)) {
          dsts += (dstId -> prop)
          return true
        }
      case None =>
    }
    for (p <- partitions) {
      if (p.update(srcId, dstId, prop)) {
        return true
      }
    }
    false
  }


  def delete(srcId: VertexId, dstId: VertexId): Boolean = {
    if (reversEdges != null) {
      reversEdges.delete(dstId, srcId)
    }
    addBuffer.get(srcId) match {
      case Some(dst) =>
        if (dst.contains(dstId)) {
          dst.remove(dstId)
          return true
        }
      case None =>
    }
    for (p <- partitions) {
      if (p.delete(srcId, dstId)) {
        return true
      }
    }
    false
  }

  def findOutVertexIds(vid: VertexId, filter: Property => Boolean = edge => true): Set[VertexId] = {
    //从partit中获取out edges
    val outIds = {
      val result = partitions.map(_.findOutVertices(vid, filter))
      if (result.isEmpty) {
        mutable.Set[VertexId]()
      } else {
        result.reduce(_ ++ _)
      }
    }
    //从 add buffer 中获取 out edges
    val addIds = addBuffer.get(vid) match {
      case Some(edges) => edges.filter(e => filter(e._2)).keys.toSet
      case None => mutable.Set[VertexId]()
    }
    (outIds ++ addIds).toSet
  }

  def findInVertexIds(vid: VertexId, filter: Property => Boolean = edge => true): Set[VertexId] = {
    if (reversEdges == null) {
      reversEdges = reverse()
    }
    reversEdges.findOutVertexIds(vid, filter)
  }

  def findOutEdges(vid: VertexId, filter: Property => Boolean = edge => true): Set[Edge] = {
    //从partit中获取out edges
    val outEdges = {
      val result = partitions.map(_.findOutEdges(vid, filter))
      if (result.isEmpty) {
        Map[VertexId, Property]()
      } else {
        result.reduce(_ ++ _)
      }
    }
    //从 add buffer 中获取 out edges
    val addEdges = addBuffer.get(vid) match {
      case Some(edges) => edges.filter(e => filter(e._2))
      case None => Map[VertexId, Property]()
    }
    val result = new mutable.HashSet[Edge]
    for ((dstId, prop) <- outEdges ++ addEdges) {
      result += Edge(vid, dstId, prop)
    }
    result.toSet
  }

  def findInEdges(vid: VertexId, filter: Property => Boolean = edge => true): Set[Edge] = {
    if (reversEdges == null) {
      reversEdges = reverse()
    }
    reversEdges.findOutEdges(vid, filter).map(edge => Edge(edge.dstId, edge.srcId, edge.prop))
  }

  def findNeighborEdges(vid: VertexId, dir: EdgeDirection,
                        filter: Property => Boolean = edge => true): Set[Edge] = {
    dir match {
      case EdgeDirection.In => findInEdges(vid, filter)
      case EdgeDirection.Out => findOutEdges(vid, filter)
      case EdgeDirection.Either =>
        val outs = findOutEdges(vid, filter)
        val ins = findInEdges(vid, filter)
        outs ++ ins
      case EdgeDirection.Both =>
        println("getNeighborEdges does not support EdgeDirection.Both. Use" +
          "EdgeDirection.Either instead.")
        Set[Edge]()
    }
  }


  private def reverseAddBuff(): mutable.Map[VertexId, mutable.Map[VertexId, Property]] = {
    val addBuffer_rev = new mutable.HashMap[VertexId, mutable.Map[VertexId, Property]]
    for ((srcId, props) <- addBuffer) {
      for ((dstId, prop) <- props) {
        addBuffer_rev.get(dstId) match {
          case Some(x) => x += (srcId -> prop)
          case None =>
            val x = mutable.Map[VertexId, Property](srcId -> prop)
            addBuffer_rev += (dstId -> x)
        }
      }
    }
    addBuffer_rev
  }

  def reverse(): LEdges = {
    val partitions_rev = partitions.map(_.reverse)
    val addBuffer_rev = reverseAddBuff()
    new LEdges(partitions_rev, addBuffer_rev, addBuffSize)
  }

  def update(): Unit = {
    if (reversEdges != null) {
      reversEdges.update()
    }
    var builder = new LEdgePartitionBuilder[Property, Property]
    for ((id1, dst) <- addBuffer) {
      for ((id2, prop) <- dst) {
        builder.add(id1, id2, prop)
      }
    }
    partitions = partitions :+ builder.toEdgePartition
    builder = null
    addBuffer.clear()
    addBuffSize = 0
  }

  def save(out: DataOutputStream) = {
    //逐条输出partitions中的边
    partitions.foreach(p => {
      p.iterator.foreach(e => GraphDBUtil.saveEdge(Edge(e.srcId, e.dstId, e.attr), out))
    })
    //逐条输出addBuffer中的边
    addBuffer.foreach(x => x._2.foreach(u => GraphDBUtil.saveEdge(Edge(x._1, u._1, u._2), out)))
  }

  def load(in: DataInputStream, schFact: SchemaFactory) = {
    reversEdges = null
    partitions = List[LEdgePartition[Property, Property]]()
    var builder = new LEdgePartitionBuilder[Property, Property]
    var e = GraphDBUtil.loadEdge(in, schFact)
    while (e != null) {
      builder.add(e.srcId, e.dstId, e.prop)
      if (builder.builderSize() >= GraphDBConf.EdgeBufferMaxSize) {
        partitions = partitions :+ builder.toEdgePartition
        builder.clear()
      }
      e = GraphDBUtil.loadEdge(in, schFact)
    }
    if (builder.builderSize() > 0) {
      partitions = partitions :+ builder.toEdgePartition
      builder = null
    }
  }

  def clear() = {
    addBuffer.clear()
    addBuffSize = 0
    partitions = List[LEdgePartition[Property, Property]]()
  }
}
